# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
import glob
import os
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO
from legged_gym import LEGGED_GYM_ROOT_DIR
DATASETS_DIR = os.path.join(LEGGED_GYM_ROOT_DIR, "datasets", "mocap_motions_go1_gazebo")
MOTION_FILES = glob.glob(os.path.join(DATASETS_DIR, "*"))



class GO1AMPCfg( LeggedRobotCfg ):

    class env( LeggedRobotCfg.env ):
        num_envs = 4096
        include_history_steps = 5  # Number of steps of history to include.
        num_observations = 45
        num_privileged_obs = 48 + 187 + 12
        reference_state_initialization = True
        reference_state_initialization_prob = 0.85
        amp_motion_files = MOTION_FILES
        ee_names = ["FL_foot", "FR_foot", "RL_foot", "RR_foot"]
        get_commands_from_joystick = False

    class init_state( LeggedRobotCfg.init_state ):
        pos = [0.0, 0.0, 0.42] # x,y,z [m]
        default_joint_angles = { # = target angles [rad] when action = 0.0
            'leg0_FL_a_hip_joint': -0.15,   # [rad]
            'leg0_FL_c_thigh_joint': 0.55,     # [rad]
            'leg0_FL_d_calf_joint': -1.5,   # [rad]

            'leg1_FR_a_hip_joint': 0.15,  # [rad]
            'leg1_FR_c_thigh_joint': 0.55,     # [rad]
            'leg1_FR_d_calf_joint': -1.5,  # [rad]

            'leg2_RL_a_hip_joint': -0.15,   # [rad]
            'leg2_RL_c_thigh_joint': 0.7,   # [rad]
            'leg2_RL_d_calf_joint': -1.5,    # [rad]

            'leg3_RR_a_hip_joint': 0.15,   # [rad]
            'leg3_RR_c_thigh_joint': 0.7,   # [rad]
            'leg3_RR_d_calf_joint': -1.5,    # [rad]
        }

    class control( LeggedRobotCfg.control ):
        # PD Drive parameters:
        control_type = 'P'
        stiffness = {'joint': 80.}  # [N*m/rad]
        damping = {'joint': 1.0}     # [N*m*s/rad]
        # action scale: target angle = actionScale * action + defaultAngle
        action_scale = 0.25
        # decimation: Number of control action updates @ sim DT per policy DT
        decimation = 4
    class hang_on_debug_angles:
        import numpy as np
        angle_list = [
            np.array([0.0,0.9,-1.8,0.0,0.9,-1.8,0.0,0.9,-1.8,0.0,0.9,-1.8],dtype=np.float32),
            np.array([0.5,0.9,-1.8,0.2,0.5,-1.5,0.5,0.9,-1.8,0.2,0.5,-1.5],dtype=np.float32),
            np.array([-0.2,0.5,-1.5,-0.5,0.9,-1.8,-0.2,0.5,-1.5,-0.5,0.9,-1.8],dtype=np.float32),
            np.array([0.0,0.6,-1.5,0.0,0.6,-1.5,0.0,0.9,-1.8,0.0,0.9,-1.8],dtype=np.float32),
            np.array([0.0,0.9,-1.8,0.0,0.9,-1.8,0.0,0.6,-1.5,0.0,0.6,-1.5],dtype=np.float32),
        ]
        angle_names = [
            "stand",
            "right",
            "left",
            "back",
            "front",
        ]

    class terrain( LeggedRobotCfg.terrain ):
        mesh_type = 'trimesh'
        measure_heights = True
        # terrain types: [smooth slope, rough slope, stairs up, stairs down, discrete]
        terrain_proportions = [0.15, 0.15, 0.3, 0.2, 0.2]

    class asset( LeggedRobotCfg.asset ):
        file = '{LEGGED_GYM_ROOT_DIR}/resources/robots/go1/urdf/go1.urdf'
        foot_name = "foot"
        penalize_contacts_on = ["thigh", "calf"]
        terminate_after_contacts_on = [
            "base"]
        self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
        robot_name = "go1"

    class domain_rand( LeggedRobotCfg.domain_rand ):
        random_flag = True
        randomize_friction = random_flag
        friction_range = [0.5, 1.25]
        randomize_restitution = random_flag #增加恢复系数
        restitution_range = [0.0, 1.0]
        randomize_base_mass = random_flag
        added_mass_range = [0., 3.]
        randomize_com_offset = random_flag
        com_offset_range = [[-0.1, 0.1], [-0.05, 0.05], [-0.03, 0.03]]
        randomize_motor_strength = random_flag
        motor_strength_range = [0.8, 1.2]
        randomize_Kp_factor = random_flag
        Kp_factor_range = [0.8, 1.2]
        randomize_Kd_factor = random_flag
        Kd_factor_range = [0.8, 1.2]
        randomize_leg_mass = random_flag
        added_leg_range = [0.8, 1.2]
        push_robots = random_flag
        push_interval_s = 15
        max_push_vel_xy = 1.0

    class noise:
        add_noise = True
        noise_level = 1.0 # scales other values
        class noise_scales:
            dof_pos = 0.03
            dof_vel = 1.5
            lin_vel = 0.1
            ang_vel = 0.3
            gravity = 0.05
            height_measurements = 0.1

    class rewards( LeggedRobotCfg.rewards ):
        soft_dof_pos_limit = 0.9
        base_height_target = 0.323
        clearance_height_target = -0.20
        class scales( LeggedRobotCfg.rewards.scales ):
            termination = 0.0
            tracking_lin_vel = 1 * 1.3
            tracking_ang_vel = 0.5 * 1.3
            lin_vel_z = -2.0
            ang_vel_xy = -0.05
            orientation = 0.0
            torques = -1e-4
            dof_vel = -0.1 * 0.005
            dof_acc = -2.5e-7
            base_height = -1.0
            feet_air_time = 1.0
            foot_clearance = -0.0
            collision = -0.15
            stumble = -0.3*2
            action_rate = 0.0
            stand_still = 0.0
            dof_pos_limits = 0.0

    class commands( LeggedRobotCfg.commands ):
        curriculum = True
        max_curriculum = 1.
        num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw, heading (in heading mode ang_vel_yaw is recomputed from heading error)
        resampling_time = 10. # time before command are changed[s]
        heading_command = True # if true: compute ang vel command from heading error
        probs_control_modes = [0.4, 0.6, 0.0]
        class ranges:
            lin_vel_x = [-1.0, 1.0] # min max [m/s]
            lin_vel_y = [-0.7, 0.7]   # min max [m/s]
            ang_vel_yaw = [-1.57, 1.57]    # min max [rad/s]
            heading = [-3.14, 3.14]

class GO1AMPCfgPPO( LeggedRobotCfgPPO ):
    runner_class_name = 'AMPOnPolicyRunner'
    class policy:
        actor_hidden_dims = [512, 256, 128]
        critic_hidden_dims = [512, 256, 128]
        hidden_size = 256
        lstm_num_layers = 3
        num_out_terrain = 16
        num_out_privilege = 8
        num_obs = 225

    class algorithm( LeggedRobotCfgPPO.algorithm ):
        entropy_coef = 0.01
        amp_replay_buffer_size = 1000000
        num_learning_epochs = 5
        num_mini_batches = 4

    class runner( LeggedRobotCfgPPO.runner ):
        run_name = ''
        experiment_name = 'go1_amp_example'
        algorithm_class_name = 'AMPPPO'
        policy_class_name = 'ActorCritic'
        student_train = True
        max_iterations = 30000 # number of policy updates
        save_interval = 500

        amp_reward_coef = 1.2 * 0.02
        amp_motion_files = MOTION_FILES
        amp_num_preload_transitions = 2000000
        amp_task_reward_lerp = 0.3
        amp_discr_hidden_dims = [1024, 512]

        min_normalized_std = [0.01, 0.01, 0.01] * 4
